package com.yololabelmaker;

import java.awt.image.BufferedImage;
import java.io.*;
import java.nio.file.*;
import java.util.*;

public class YOLOLabelLoader {
    // 类别映射表，用于存储类别名称和ID的对应关系
    private static java.util.Map<Integer, String> classMap = new java.util.HashMap<>();
    
    static {
        // 尝试加载已有的类别映射文件
        loadClassMap();
    }
    
    // 加载标签
    public List<LabelRect> loadLabels(String imagePath) {
        if (imagePath == null) {
            return null;
        }
        
        // 找到对应的标签文件
        File imageFile = new File(imagePath);
        String labelsDirPath;
        
        // 检查是否是网络图片
        if (imagePath.startsWith("network_image_")) {
            // 对于网络图片，从当前目录的labels文件夹加载
            labelsDirPath = System.getProperty("user.dir") + File.separator + "labels";
        } else {
            // 对于本地图片，按照YOLO数据集结构查找
            String parentDir = imageFile.getParent();
            if (parentDir != null && parentDir.endsWith("images")) {
                // 如果图片在images文件夹中，则标签应该放在同级的labels文件夹中
                labelsDirPath = parentDir.replace("images", "labels");
            } else {
                // 否则，在图片所在目录的labels文件夹中查找
                labelsDirPath = parentDir + File.separator + "labels";
            }
        }
        
        // 构建标签文件路径
        String imageFileName = imageFile.getName();
        String labelFileName = imageFileName.substring(0, imageFileName.lastIndexOf('.')) + ".txt";
        String labelFilePath = labelsDirPath + File.separator + labelFileName;
        
        File labelFile = new File(labelFilePath);
        if (!labelFile.exists()) {
            // 尝试在图片同一目录下查找标签文件
            labelFilePath = imageFile.getParent() + File.separator + labelFileName;
            labelFile = new File(labelFilePath);
            if (!labelFile.exists()) {
                return null; // 没有找到标签文件
            }
        }
        
        List<LabelRect> labelRects = new ArrayList<>();
        
        try (BufferedReader reader = new BufferedReader(new FileReader(labelFile))) {
            String line;
            while ((line = reader.readLine()) != null) {
                line = line.trim();
                if (line.isEmpty() || line.startsWith("#")) {
                    continue; // 跳过空行和注释行
                }
                
                // 解析YOLO格式的标签行
                LabelRect rect = parseYOLOLine(line);
                if (rect != null) {
                    labelRects.add(rect);
                }
            }
        } catch (IOException e) {
            System.err.println("加载标签文件时发生错误: " + e.getMessage());
            return null;
        }
        
        return labelRects;
    }
    
    // 解析YOLO格式的标签行
    private LabelRect parseYOLOLine(String line) {
        try {
            String[] parts = line.split("\\s+");
            if (parts.length < 5) {
                return null;
            }
            
            int classId = Integer.parseInt(parts[0]);
            double centerX = Double.parseDouble(parts[1]);
            double centerY = Double.parseDouble(parts[2]);
            double width = Double.parseDouble(parts[3]);
            double height = Double.parseDouble(parts[4]);
            
            // 创建标签矩形对象
            LabelRect rect = new LabelRect();
            rect.setClassId(classId);
            
            // 从类别映射中获取类别名称
            String className = getClassName(classId);
            if (className != null) {
                rect.setClassName(className);
            } else {
                rect.setClassName("class_" + classId);
            }
            
            // 保存原始YOLO坐标比例信息，供后续使用
            // 实际的像素坐标转换将在图片加载后进行
            rect.setUserData(centerX, centerY, width, height);
            
            return rect;
        } catch (NumberFormatException e) {
            System.err.println("解析标签行时发生错误: " + line);
            return null;
        }
    }
    
    // 获取类别名称
    private String getClassName(int classId) {
        return classMap.get(classId);
    }
    
    // 加载类别映射文件
    private static void loadClassMap() {
        String classMapFilePath = System.getProperty("user.dir") + File.separator + "classes.txt";
        File classMapFile = new File(classMapFilePath);
        if (classMapFile.exists()) {
            try (BufferedReader reader = new BufferedReader(new FileReader(classMapFile))) {
                String line;
                int id = 0;
                while ((line = reader.readLine()) != null) {
                    line = line.trim();
                    if (!line.isEmpty()) {
                        classMap.put(id++, line);
                    }
                }
            } catch (IOException e) {
                System.err.println("加载类别映射文件时发生错误: " + e.getMessage());
            }
        }
    }
    
    // 更新标签框的坐标（在图片尺寸已知的情况下）
    public static void updateLabelRectCoordinates(LabelRect rect, int imageWidth, int imageHeight, 
                                                double centerX, double centerY, 
                                                double width, double height) {
        // 计算实际像素坐标
        int pixelCenterX = (int) (centerX * imageWidth);
        int pixelCenterY = (int) (centerY * imageHeight);
        int pixelWidth = (int) (width * imageWidth);
        int pixelHeight = (int) (height * imageHeight);
        
        // 设置标签框的坐标和尺寸
        rect.setX(pixelCenterX - pixelWidth / 2);
        rect.setY(pixelCenterY - pixelHeight / 2);
        rect.setWidth(pixelWidth);
        rect.setHeight(pixelHeight);
    }
}